#coding=utf-8
from caffe2.python import workspace, model_helper
import numpy as np

def main():
    # 创建输入数据
    data    = np.random.rand(16,100).astype(np.float32)
    # 创建标注信息, 整数[0, 9]
    label   = (np.random.rand(16)*10).astype(np.int32)

    # 将数据和标签喂进去
    workspace.FeedBlob("data", data)
    workspace.FeedBlob("label", label)

    # 采用model_helper 创建模型
    m   = model_helper.ModelHelper(name="my_1st_net")

    fc_in_num   = 100
    fc_out_num  = 10
    minisize    = 16
    # 创建FC的权重
    FC_weight   = m.param_init_net.XavierFill([], 'fc_w', shape=[fc_out_num, fc_in_num])
    FC_bias     = m.param_init_net.ConstantFill([], 'fc_b', shape=[fc_out_num, ])

    # 构建layer, 输出, loss
    fc_1    =   m.net.FC(["data", 'fc_w', 'fc_b'], 'fc1')
    pred    =   m.net.Sigmoid(fc_1, 'pred')
    # 下面的两种softmax都可以运行
    #softmax, loss   = m.net.SoftmaxWithLoss(["pred", "label"], ["softmax", "loss"])
    softmax, loss   = m.net.SoftmaxWithLoss([pred, "label"], ["softmax", "loss"])
    print("add gradient ops for each op in forward pass")
    m.AddGradientOperators([loss])
    print("(m.net.Proto()")
    print(m.net.Proto())
    print("m.param_init_net.Proto()")
    print(m.param_init_net.Proto())
    print("param init")
    workspace.RunNetOnce(m.param_init_net)
    print("create actual net in workspace")
    workspace.CreateNet(m.net)
    print("forward it")
    # run 100 x 10 iterations
    for _ in range(100):
        data    = np.random.rand(minisize,fc_in_num).astype(np.float32)
        label   = (np.random.rand(minisize)*fc_out_num).astype(np.int32)
        workspace.FeedBlob("data", data)
        workspace.FeedBlob("label", label)
        workspace.RunNet(m.name, 10)
    print("Softmax:{}, Loss:{}".format(workspace.FetchBlob("softmax"), workspace.FetchBlob("loss")))

if __name__ == "__main__":
    main()
